import time
from typing import Union, Callable

from scripts.actions.Context import Context
from scripts.config.Config import Config
from scripts.device.AdbDevice import AdbDevice
from scripts.Log import Log
from scripts.actions.Action import Action
from scripts.util.ProgressHelper import ProgressHelper
from scripts.util.RebootHelper import RebootHelper


class PrepareAction(Action):
    KEY = Action.register('prepare')
    DISPLAY_NAME = 'Prepare'

    NO_SAFE_REBOOT = Config.STEP_PREPARE_NO_SAFE_REBOOT

    def __init__(self, context: Context):
        super().__init__(PrepareAction.KEY, PrepareAction.DISPLAY_NAME, context, False)

    def _on_backup(self, progress: Union[None, Callable[[int, int, str], None]]):
        self._on_prepare(progress)

    def _on_restore(self, progress: Union[None, Callable[[int, int, str], None]]):
        self._on_prepare(progress)

    def _on_prepare(self, progress: Union[None, Callable[[int, int, str], None]]) -> Union[None, AdbDevice]:
        prog = ProgressHelper(progress, 1)
        prog.on_message('Please plug in device [%s] and enable adb!!' % self.sn())

        Log.d('kill previous adb server')
        AdbDevice.kill_server()

        # get device infos
        device = self.wait_for_device()
        with device.shell() as shell:
            self.set_android_version(shell.android_version())
            self.set_device_model(shell.device_model())

        # print device info
        Log.i('=========== Device Info ============')
        Log.i('> android version: %s' % self.android_version())
        Log.i('> device model: %s' % self.device_model())
        Log.i('====================================')

        # root
        device = self.wait_for_device_rooted()
        if device is None:
            Log.d('adb root failed. do device root')
            device = self.wait_for_device()
            with device.shell() as shell:
                root_success: bool
                ver = self.android_version()
                Log.d('android version: %d' % ver)
                if ver == 9:
                    root_success = self.root9(device)
                elif ver == 11:
                    root_success = self.root11(device)
                else:
                    Log.e('unable to root device with android version: %d' % ver)
                    root_success = False

                # check root result
                if not root_success:
                    Log.e('device root failed')
                    prog.inc().on_message('Done [%s]' % self.duration_time_str())
                    return None

            # adb root & reboot in case device require reboot for disable verity
            Log.d('wait for device: %s' % self.sn())
            device = self.wait_for_device_rooted()
            Log.d('reboot again for disable verity')
            device.reboot()

            # adb root again
            Log.d('wait for device: %s' % self.sn())
            device = self.wait_for_device_rooted()

            Log.d('device rooted!!')
        else:
            Log.d('unnecessary to root device')

        # remove screen lock and safe reboot to save current changes
        with device.shell() as shell:
            # ask user to remove screen lock
            users = shell.pm_list_users()
            for user in users:
                while not shell.locksettings_get_disabled(user.user_id):
                    prog.on_message('!!Please remove screen lock of user: %s' % user.user_name)
                    Log.d('ask user to remove screen lock of user: %s' % user.user_name)
                    time.sleep(1)

            # safe reboot to apply system changes
            prog.on_message('Reboot to save all system changes')
            if PrepareAction.NO_SAFE_REBOOT:
                Log.d('ignore safe reboot due to NO_SAFE_REBOOT is True')
            else:
                Log.d('safe reboot to apply system changes')
                _, device, shell = RebootHelper(self).safe_reboot(device, shell)

        # create temp dir
        device = self.wait_for_boot_up_ready()
        with device.shell() as shell:
            shell.mkdir_and_check(self.device_temp_dir())

        prog.inc().on_message('Done [%s]' % self.duration_time_str())

        return device

    def root9(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.enable_root_9()

        # reboot
        fastboot_device.reboot()

        return result

    def root11(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.enable_root_11()

        # reboot
        fastboot_device.reboot()

        return result
